#include "Arduino.h"
#include "Quaternion.h"
#include <FP_math.h>

  
Quaternion::Quaternion(void) {
  setQuat(1, 0, 0, 0);
}
  
Quaternion::Quaternion(int q[]) {
  setQuat(q);
}
  
Quaternion::Quaternion(int h, int p, int r) {
  setQuat(h, p, r);
}
  
Quaternion::Quaternion(int q1, int q2, int q3, int q4) {
  setQuat(q1, q2, q3, q4);
}

void Quaternion::setQuat(int q[]) {
  w = q[0];
  x = q[1];
  y = q[2];
  z = q[3];
  normalize();
  rSet = false;
}
  
void Quaternion::setQuat(int q1, int q2, int q3, int q4) {
  w = q1;
  x = q2;
  y = q3;
  z = q4;
  normalize();
  rSet = false;
}

  
void Quaternion::setQuat(int h, int p, int r) {

  // Initializes Quaternion to (h)eading, (p)itch, (r)oll 

  int cos_h2 = FP_math::fp_cos_2(h);
  int sin_h2 = FP_math::fp_sin_2(h);
  int cos_p2 = FP_math::fp_cos_2(p);
  int sin_p2 = FP_math::fp_sin_2(p);
  int cos_r2 = FP_math::fp_cos_2(r);
  int sin_r2 = FP_math::fp_sin_2(r);
  
  w =  FP_math::mult16_14(FP_math::mult16_14(cos_p2, cos_r2), cos_h2);
  w += FP_math::mult16_14(FP_math::mult16_14(sin_h2, sin_r2), sin_p2);

  x =  FP_math::mult16_14(FP_math::mult16_14(cos_h2, cos_p2), sin_r2);
  x -= FP_math::mult16_14(FP_math::mult16_14(cos_r2, sin_h2), sin_p2);

  y =  FP_math::mult16_14(FP_math::mult16_14(cos_h2, cos_r2), sin_p2);
  y += FP_math::mult16_14(FP_math::mult16_14(sin_h2, cos_p2), sin_r2);

  z =  FP_math::mult16_14(FP_math::mult16_14(-cos_h2, sin_p2), sin_r2);
  z += FP_math::mult16_14(FP_math::mult16_14( cos_r2, cos_p2), sin_h2);
  
  normalize();
  
  rSet = false;
}

int Quaternion::getW(void){
  return w;
}

int Quaternion::getX(void){
  return x;
}

int Quaternion::getY(void){
  return y;
}

int Quaternion::getZ(void){
  return z;
}

void Quaternion::getQ(int q[]){
  q[0] = w;
  q[1] = x;
  q[2] = y;
  q[3] = z;
}
  
void Quaternion::multiply(Quaternion a) {
  multiply(a.w, a.x, a.y, a.z);
}
  
void Quaternion::multiply(int q[]) {
  multiply(q[0], q[1], q[2], q[3]);
}
  
void Quaternion::multiply(int theW, int theX, int theY, int theZ) {
  int tempW = FP_math::mult16_14(theW, w) - FP_math::mult16_14(theX, x);
  tempW +=   -FP_math::mult16_14(theY, y) - FP_math::mult16_14(theZ, z);

  int tempX = FP_math::mult16_14(theW, x) + FP_math::mult16_14(theX, w);
  tempX +=    FP_math::mult16_14(theY, z) - FP_math::mult16_14(theZ, y);

  int tempY = FP_math::mult16_14(theW, y) - FP_math::mult16_14(theX, z);
  tempY +=    FP_math::mult16_14(theY, w) + FP_math::mult16_14(theZ, x);

  int tempZ = FP_math::mult16_14(theW, z) + FP_math::mult16_14(theX, y);
  tempZ +=   -FP_math::mult16_14(theY, x) + FP_math::mult16_14(theZ, w);

  w = tempW;
  x = tempX;
  y = tempY;
  z = tempZ;

  rSet = false;

}
  
void Quaternion::postMultiply(Quaternion a) {
  postMultiply(a.w, a.x, a.y, a.z);
}
  
void Quaternion::postMultiply(int q[]) {
  postMultiply(q[0], q[1], q[2], q[3]);
}
  
void Quaternion::postMultiply(int theW, int theX, int theY, int theZ) {
  int tempW = FP_math::mult16_14(w, theW) - FP_math::mult16_14(x, theX);
  tempW +=   -FP_math::mult16_14(y, theY) - FP_math::mult16_14(z, theZ);

  int tempX = FP_math::mult16_14(w, theX) + FP_math::mult16_14(x, theW);
  tempX +=    FP_math::mult16_14(y, theZ) - FP_math::mult16_14(z, theY);

  int tempY = FP_math::mult16_14(w, theY) - FP_math::mult16_14(x, theZ);
  tempY +=    FP_math::mult16_14(y, theW) + FP_math::mult16_14(z, theX);

  int tempZ = FP_math::mult16_14(w, theZ) + FP_math::mult16_14(x, theY);
  tempZ +=   -FP_math::mult16_14(y, theX) + FP_math::mult16_14(z, theW);

  w = tempW;
  x = tempX;
  y = tempY;
  z = tempZ;

  rSet = false;

}
  
void Quaternion::negMult(Quaternion a) {
  negMult(a.w, a.x, a.y, a.z);
}
  
void Quaternion::negMult(int q[]) {
  negMult(q[0], q[1], q[2], q[3]);
}
  
void Quaternion::negMult(int theW, int theX, int theY, int theZ) {
  int tempW = FP_math::mult16_14(theW, w) + FP_math::mult16_14(theX, x);
  tempW +=    FP_math::mult16_14(theY, y) + FP_math::mult16_14(theZ, z);

  int tempX = FP_math::mult16_14(theW, x) - FP_math::mult16_14(theX, w);
  tempX +=   -FP_math::mult16_14(theY, z) + FP_math::mult16_14(theZ, y);

  int tempY = FP_math::mult16_14(theW, y) + FP_math::mult16_14(theX, z);
  tempY +=   -FP_math::mult16_14(theY, w) - FP_math::mult16_14(theZ, x);

  int tempZ = FP_math::mult16_14(theW, z) - FP_math::mult16_14(theX, y);
  tempZ +=    FP_math::mult16_14(theY, x) - FP_math::mult16_14(theZ, w);

  w = tempW;
  x = tempX;
  y = tempY;
  z = tempZ;

  rSet = false;

}
  
  
  
void Quaternion::rotationMatrix(int RMx[][3]) {
  if (!rSet) setRotationMatrix();
  for (int r = 0; r < 3; r++) {
    for (int c = 0; c < 3; c++) {
      RMx[r][c] = R[r][c];
    }
  }
}

void Quaternion::setRotationMatrix(void) {
  int xx = FP_math::mult16_14(x, x) << 1;
  int yy = FP_math::mult16_14(y, y) << 1;
  int zz = FP_math::mult16_14(z, z) << 1;
  int xy = FP_math::mult16_14(x, y) << 1;
  int xz = FP_math::mult16_14(x, z) << 1;
  int xw = FP_math::mult16_14(x, w) << 1;
  int yz = FP_math::mult16_14(y, z) << 1;
  int yw = FP_math::mult16_14(y, w) << 1;
  int zw = FP_math::mult16_14(z, w) << 1;
  R[0][0] = FP_math::oneFix14 - yy - zz;
  R[0][1] = xy - zw;
  R[0][2] = xz + yw;
  R[1][0] = xy + zw;
  R[1][1] = FP_math::oneFix14 - xx - zz;
  R[1][2] = yz - xw;
  R[2][0] = xz - yw;
  R[2][1] = yz + xw;
  R[2][2] = FP_math::oneFix14 - xx - yy;

  rSet = true;
}

void Quaternion::body2ned(int vector[], int outVector[]) {
  if (!rSet) setRotationMatrix();
  FP_math::matMult(vector, R, outVector);
}
  
void Quaternion::ned2body(int vector[], int outVector[]) {
  if (!rSet) setRotationMatrix();
  FP_math::matMultT(vector, R, outVector);
}

void Quaternion::body2nedXY(int vector[], int outVector[]) {
  if (!rSet) setRotationMatrix();
  for (int r = 0; r < 2; r++) {
    outVector[r] = 0;
    for (int c = 0; c < 2; c++) {
      outVector[r] += FP_math::mult16_14(R[r][c], vector[c]);
    }
  }
}

int Quaternion::scaleZ(void) {
  if (!rSet) setRotationMatrix();
  return R[2][2];
}

  
void Quaternion::ned2bodyXY(int vector[], int outVector[]) {
  if (!rSet) setRotationMatrix();
  for (int r = 0; r < 2; r++) {
    outVector[r] = 0;
    for (int c = 0; c < 2; c++) {
      outVector[r] += FP_math::mult16_14(R[c][r], vector[c]);
    }
  }
}

void Quaternion::normalize() {
  int norm = FP_math::mult16_14(w, w) + FP_math::mult16_14(x, x);
  norm +=    FP_math::mult16_14(y, y) + FP_math::mult16_14(z, z);
  norm =     FP_math::div16_14(FP_math::oneFix14, FP_math::fp_sqrt(norm));

  w = FP_math::mult16_14(w, norm);
  x = FP_math::mult16_14(x, norm);
  y = FP_math::mult16_14(y, norm);
  z = FP_math::mult16_14(z, norm);
  
  rSet = false;
}
  

// ********************************************************
// DifQuat Sub-class Implementation

  
DifQuat::DifQuat() {
  zeroVals();
}
  
DifQuat::DifQuat(int q[]) : Quaternion(q) {
  zeroVals();
}
  
DifQuat::DifQuat(int h, int p, int r) : Quaternion(h, p, r) {
  zeroVals();
}
  
DifQuat::DifQuat(int q1, int q2, int q3, int q4) : Quaternion(q1, q2, q3, q4) {
  zeroVals();
}

void DifQuat::zeroVals(void) {
  dx = 0;
  dy = 0;
  dz = 0;
}


void DifQuat::computeDiff(Quaternion a) {
  computeDiff(a.getW(), a.getX(), a.getY(), a.getZ());
}
  
void DifQuat::computeDiff(int q[]) {
  computeDiff(q[0], q[1], q[2], q[3]);
}
  
void DifQuat::computeDiff(int theW, int theX, int theY, int theZ) {
  
  // Computes the body angle difference from input to Quaternion
  // Difference Angles are in Radians Fix_16_14

  int tempW = w;
  int tempX = x;
  int tempY = y;
  int tempZ = z;
  negMult(theW, theX, theY, theZ);
  dx = x << 1;
  dy = y << 1;
  dz = z << 1;
  w = tempW;
  x = tempX;
  y = tempY;
  z = tempZ;
}

// ********************************************************
// IMU Sub-class Implementation

IMU::IMU() {
  load_int();
}

IMU::IMU(int q[]) : Quaternion(q) {
  load_int();
}
  
IMU::IMU(int h, int p, int r) : Quaternion(h, p, r) {
  load_int();
}
  
IMU::IMU(int q1, int q2, int q3, int q4) : Quaternion(q1, q2, q3, q4) {
  load_int();
}

void IMU::load_int(void) {

  // Initializes IMU integrators

  int n_shift = 16;
  w_int.l = long(w) << n_shift;
  x_int.l = long(x) << n_shift;
  y_int.l = long(y) << n_shift;
  z_int.l = long(z) << n_shift;
}

void IMU::setQuat(int q[]) {
  Quaternion::setQuat(q);
  load_int();
}
  
void IMU::setQuat(int q1, int q2, int q3, int q4) {
  Quaternion::setQuat(q1, q2, q3, q4);
  load_int();
}

  
void IMU::setQuat(int h, int p, int r) {
  Quaternion::setQuat(h, p, r);
  load_int();
}
  


int IMU::heading(void) {

  // Computes heading angle (deg10) of Quaternion 

  int yTan = (FP_math::mult16_14(x, y) + FP_math::mult16_14(z, w)) << 1;
  int xTan = FP_math::oneFix14 - ((FP_math::mult16_14(y, y) + FP_math::mult16_14(z, z)) << 1);
  int h = FP_math::fp_atan2(yTan, xTan);
  if (h < 0) h += FP_math::D360;
  return h;
}
  

int IMU::pitch(void) {

// Computes pitch angle (deg10) of Quaternion

  int a = (FP_math::mult16_14(y, w) - FP_math::mult16_14(x, z)) << 1;
  return FP_math::fp_asin(a);
}
  

int IMU::roll(void) {

  // Computes roll angle (deg10) of Quaternion

  int yTan = (FP_math::mult16_14(y, z) + FP_math::mult16_14(x, w)) << 1;
  int xTan = FP_math::oneFix14 - ((FP_math::mult16_14(x, x) + FP_math::mult16_14(y, y)) << 1);
  return FP_math::fp_atan2(yTan, xTan);
}


void IMU::propagate(int xDot, int yDot, int zDot, int dT) {

  // Updates Quaternion based on Gyro Body Rates
  // xDot, yDot, zDot are Raw Integer Rates from gyro
  // dT is in milliseconds
  // Gyro Scale Factor is set through setGyroScale() method
  // setGyroScale() takes care of integrator Fix_32_30 scaling and
  // millisecond - second - 1/1024s scaling

  // The Quaternion multiply:
  // Ref: setGyroScale for Shift parameter explanations

  long temp = long(-FP_math::mult16_14(x, xDot));
  temp -= FP_math::mult16_14(y, yDot);
  temp -= FP_math::mult16_14(z, zDot);
  w_int.l += temp * gyroScale * dT;

  temp = long(FP_math::mult16_14(w, xDot));
  temp -= FP_math::mult16_14(z, yDot);
  temp += FP_math::mult16_14(y, zDot);
  x_int.l += temp * gyroScale * dT;

  temp = long(FP_math::mult16_14(z, xDot));
  temp += FP_math::mult16_14(w, yDot);
  temp -= FP_math::mult16_14(x, zDot);
  y_int.l += temp * gyroScale * dT;

  temp = long(-FP_math::mult16_14(y, xDot));
  temp += FP_math::mult16_14(x, yDot);
  temp += FP_math::mult16_14(w, zDot);
  z_int.l += temp * gyroScale * dT;
  
  int norm =  FP_math::mult16_14(w_int.s.hi, w_int.s.hi) + FP_math::mult16_14(x_int.s.hi, x_int.s.hi);
  norm +=     FP_math::mult16_14(y_int.s.hi, y_int.s.hi) + FP_math::mult16_14(z_int.s.hi, z_int.s.hi);
  norm -= FP_math::oneFix14;
  norm = norm >> 1;

  w_int.l -= normOne(w_int.s.hi, norm);
  x_int.l -= normOne(x_int.s.hi, norm);
  y_int.l -= normOne(y_int.s.hi, norm);
  z_int.l -= normOne(z_int.s.hi, norm);

  w = w_int.s.hi;
  x = x_int.s.hi;
  y = y_int.s.hi;
  z = z_int.s.hi;

  rSet = false;
}

void IMU::propagate(int vector[], int dT) {
  propagate(vector[0], vector[1], vector[2], dT);
}

void IMU::propagate(Vector3 vector, int dT) {
  propagate(vector.s.x, vector.s.y, vector.s.z, dT);
}

long IMU::normOne(int val, int norm) {

// Computes (Fix_32_30) temp = (Fix_16_14) val * (Fix_16_14) norm

  long temp = long(val);
  temp *= norm;
  temp = temp << 2;
  return temp;
}


void IMU::setGyroScale(int scale) {

  // Values scale Gyro Raw Int Value to Radians/s
  // Private properties set to accommodate internal Fix_32_30 format
  // of internal integrators & dT in milliseconds
  //
  // Rate (Radians/sec) = Gyro_Raw * scale * (2^-exp2)

  // This scales mantissa by 1028/1000 to accommodate (>> 10) ~ (1/1000)
  long temp = long(scale) << 10;
  temp /= 1000;
  gyroScale = int(temp);

}

